import cv2
import numpy as np
from matplotlib import pyplot as plot


def make_transform_matrix(image, d):
    transfor_matrix = np.zeros(image.shape)
    center_point = tuple(map(lambda x: (x - 1) / 2, image.shape))
    for i in range(transfor_matrix.shape[0]):
        for j in range(transfor_matrix.shape[1]):
            def cal_distance(pa, pb):
                from math import sqrt
                dis = sqrt((pa[0] - pb[0]) ** 2 + (pa[1] - pb[1]) ** 2)
                return dis

            dis = cal_distance(center_point, (i, j))
            transfor_matrix[i, j] = 1 - np.exp(-(dis ** 2) / (2 * (d ** 2)))
    return transfor_matrix

def make_transform_matrix_d(image, d):
    transfor_matrix = np.zeros(image.shape)
    center_point = tuple(map(lambda x: (x - 1) / 2, image.shape))
    for i in range(transfor_matrix.shape[0]):
        for j in range(transfor_matrix.shape[1]):
            def cal_distance(pa, pb):
                from math import sqrt
                dis = sqrt((pa[0] - pb[0]) ** 2 + (pa[1] - pb[1]) ** 2)
                return dis

            dis = cal_distance(center_point, (i, j))
            transfor_matrix[i, j] = np.exp(-(dis ** 2) / (2 * (d ** 2)))
    return transfor_matrix

def make_transform_matrix_b(image, d, n):
    transfor_matrix = np.zeros(image.shape)
    center_point = tuple(map(lambda x: (x - 1) / 2, image.shape))
    for i in range(transfor_matrix.shape[0]):
        for j in range(transfor_matrix.shape[1]):
            def cal_distance(pa, pb):
                from math import sqrt
                dis = sqrt((pa[0] - pb[0]) ** 2 + (pa[1] - pb[1]) ** 2)
                return dis

            dis = cal_distance(center_point, (i, j))
            transfor_matrix[i, j] = 1 / ((1 + (d / dis)) ** n)
    return transfor_matrix

img = src = cv2.imread('lena.bmp', cv2.IMREAD_COLOR)
src = cv2.cvtColor(src, cv2.COLOR_BGR2RGB)
h, img, s = cv2.split(cv2.cvtColor(src, cv2.COLOR_RGB2HLS_FULL))
f = np.fft.fft2(img)
fshift = np.fft.fftshift(f)
magnitude_spectrum = 20 * np.log(np.abs(fshift))

d_matrix = make_transform_matrix(img, 20)
fshift1 = fshift*d_matrix
magnitude_spectrum1 = 20 * np.log(np.abs(fshift1))

f_ishift = np.fft.ifftshift(fshift1)
img_back = np.fft.ifft2(f_ishift)
img_back = np.abs(img_back)

d_matrix2 = make_transform_matrix_d(img, 30)
fshift2 = fshift*d_matrix2
magnitude_spectrum2 = 20 * np.log(np.abs(fshift2))

f_ishift2 = np.fft.ifftshift(fshift2)
img_back2 = np.fft.ifft2(f_ishift2)
img_back2 = np.abs(img_back2)
img_back2[img_back2 > 255] = 255
img_back2[img_back2 < 0] = 0
img_back2 = img_back2.astype(np.uint8)
img_back2 = cv2.merge((h, img_back2, s))
img_back2 = cv2.cvtColor(img_back2, cv2.COLOR_HLS2RGB_FULL)

plot.subplot(231), plot.imshow(src)
plot.title("Input"), plot.xticks([]), plot.yticks([])

plot.subplot(234), plot.imshow(magnitude_spectrum, cmap="bwr")
plot.title('Spectrum'), plot.xticks([]), plot.yticks([])

plot.subplot(232), plot.imshow(img_back, cmap='gray')
plot.title("Gauss High Filter"), plot.xticks([]), plot.yticks([])

plot.subplot(235), plot.imshow(magnitude_spectrum1, cmap = "bwr")
plot.title('High Spectrum'), plot.xticks([]), plot.yticks([])

plot.subplot(233), plot.imshow(img_back2)
plot.title("Gauss Low Filter"), plot.xticks([]), plot.yticks([])

plot.subplot(236), plot.imshow(magnitude_spectrum2, cmap = "bwr")
plot.title('Low Spectrum'), plot.xticks([]), plot.yticks([])

plot.show()
